#pragma once

#include <mutex>
#include <condition_variable>
#include <cassert>
#include <atomic>

namespace kvalloc {

class SpinLock {
public:
    std::atomic_bool lock_;
public:
    SpinLock() {
        lock_.store(false);
    }

    ~SpinLock() {
        lock_.store(false);
    }

    void Lock(){
        while(lock_.exchange(true)){}
    }

    bool TryLock() {
        return !lock_.exchange(true);
    }

    void UnLock(){
        while(!lock_.exchange(false)){}
    }

    void Wait(){
        while(lock_.load()){};
    }
};

class MutexLock {
 public:
  explicit MutexLock(std::mutex* mu): mu_(mu) {
    this->mu_->lock();
  }
  ~MutexLock() { 
    this->mu_->unlock(); 
  }

  MutexLock(const MutexLock&) = delete;
  MutexLock& operator=(const MutexLock&) = delete;

private:
  std::mutex * const mu_;
};

// Thinly wraps std::condition_variable.
class CondVar {
public:
  explicit CondVar(std::mutex * mu) : mu_(mu) { 
    assert(mu != nullptr); 
  }
  ~CondVar() = default;

  CondVar(const CondVar&) = delete;
  CondVar& operator=(const CondVar&) = delete;

  void Wait() {
    std::unique_lock<std::mutex> lock(*mu_, std::adopt_lock);
    cv_.wait(lock);
    lock.release();
  }
  void Signal() { cv_.notify_one(); }
  void SignalAll() { cv_.notify_all(); }

 private:
  std::condition_variable cv_;
  std::mutex* const mu_;
};

// A compact reader writer lock: no more than 63 readers
// readers will wait for on-going writers
// writers will wait for on-going readers
class RWLock {
private:
    uint8_t lock : 1;  // exclusive bits for 
    uint8_t reader_num : 7;

public:
    RWLock() {
        lock = 0;
        reader_num = 0;
    }

    void RLock() { // atomically add the reader_num
		RWLock expected = *this;
		RWLock desired;

        bool success;
        do {
            // expected = *this;
            expected.lock = 0;   // if locked, wait
            desired.reader_num = expected.reader_num + 1; // add one
            success = __atomic_compare_exchange(this, &expected, &desired, false, __ATOMIC_ACQUIRE, __ATOMIC_ACQUIRE);
        } while(success == false);
    }

    void WLock() { // atomically set the intent and lock bit, make sure reader_num == 0
		RWLock expected;
        RWLock desired;
		desired.lock = 1;   // desire to be locked
        expected.reader_num = this->reader_num;
    
        bool success;
        do {
            // expected = *this;
            expected.lock = 0;       // if locked, wait
            desired.reader_num = expected.reader_num; // do not update reader_num
            success = __atomic_compare_exchange(this, &expected, &desired, false, __ATOMIC_ACQUIRE, __ATOMIC_ACQUIRE);
        } while(success == false);

        RWLock tmp;
        do {
            __atomic_load(this, &tmp, __ATOMIC_ACQUIRE);
        } while(tmp.reader_num > 0);
    }

    void UnRLock() { // atomically sub the reader_num
		RWLock expected = *this;
        RWLock desired;  // all zero

        bool success;
        do {
            // expected = *this;
            desired.lock = expected.lock; // locked or not, whatever
            desired.reader_num = expected.reader_num - 1; // sub one
            success = __atomic_compare_exchange(this, &expected, &desired, false, __ATOMIC_ACQUIRE, __ATOMIC_ACQUIRE);
        } while(success == false);
    }

    void UnWLock() { // clear all bits
        RWLock desired;  // all zero
        RWLock old;
        __atomic_exchange(this, &desired, &old, __ATOMIC_ACQUIRE);
    }

    void Clear() {
        lock = 0;
        reader_num = 0;
    }
};

// void test_reader(int i) {
//     lock.RLock();
//     sleep(1);
//     printf("reader %d\n", i);
//     lock.UnRLock();
// }

// void test_writer(int i) {
//     lock.WLock();
//     sleep(1);
//     printf("writers %d\n", i);
//     lock.UnWLock();
// }

// int main() {
//     const int READER_NUM = 8;
//     const int WRITER_NUM = 8;

//     std::thread readers[READER_NUM];
//     for(int i = 0; i < READER_NUM; i++) {
//         readers[i] = std::move(std::thread(test_reader, i));
//     }

//     std::thread writers[WRITER_NUM];
//     for(int i = 0; i < WRITER_NUM; i++) {
//         writers[i] = std::move(std::thread(test_writer, i));
//     }

//     for(int i = 0; i < READER_NUM; i++) 
//         readers[i].join();
    
//     for(int i = 0; i < WRITER_NUM; i++)
//         writers[i].join(); 

//     return 0;
// }
}